package com.github.xuejike.springboot.generator;

import com.baomidou.mybatisplus.annotation.IdType;
import com.baomidou.mybatisplus.annotation.TableId;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.generator.AutoGenerator;
import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
import com.baomidou.mybatisplus.generator.config.po.TableField;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.config.rules.IColumnType;
import org.hibernate.annotations.SQLDelete;
import org.hibernate.annotations.Where;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.persistence.*;
import java.io.Serializable;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.BiFunction;


public class JpaAutoGenerator extends AutoGenerator {
    private static final Logger logger = LoggerFactory.getLogger(JpaAutoGenerator.class);
    protected BiFunction<String,String,IColumnType> fieldTypeHandler;
    protected BiFunction<String,String,IColumnType> fieldValueGenerationHandler;



    @Override
    protected ConfigBuilder pretreatmentConfigBuilder(ConfigBuilder config) {
        /**
         * 注入自定义配置
         */
        if (null != injectionConfig) {
            injectionConfig.initMap();
            config.setInjectionConfig(injectionConfig);
        }
        /**
         * 表信息列表
         */
        List<TableInfo> tableList = this.getAllTableInfoList(config);
        for (TableInfo tableInfo : tableList) {
            /* ---------- 添加导入包 ---------- */
            tableInfo.setImportPackages(Entity.class.getCanonicalName());
            tableInfo.getImportPackages().remove(com.baomidou.mybatisplus.annotation.TableField.class.getCanonicalName());
            tableInfo.getImportPackages().remove(TableId.class.getCanonicalName());
            tableInfo.getImportPackages().remove(IdType.class.getCanonicalName());
            if (config.getGlobalConfig().isActiveRecord()) {
                // 开启 ActiveRecord 模式
                logger.error("不支持 ActiveRecord 模式");
            }
            if (tableInfo.isConvert()) {
                // 表注解
                tableInfo.setImportPackages(Table.class.getCanonicalName());
                tableInfo.setImportPackages(Column.class.getCanonicalName());
            }
            if (config.getStrategyConfig().getLogicDeleteFieldName() != null && tableInfo.isLogicDelete(config.getStrategyConfig().getLogicDeleteFieldName())) {
                // 逻辑删除注解
                tableInfo.setImportPackages(SQLDelete.class.getCanonicalName());
                tableInfo.setImportPackages(Where.class.getCanonicalName());
            }
            if (StringUtils.isNotEmpty(config.getStrategyConfig().getVersionFieldName())) {
                // 乐观锁注解
                tableInfo.setImportPackages(Version.class.getCanonicalName());
            }
            if (StringUtils.isNotEmpty(config.getSuperEntityClass())) {
                // 父实体
                tableInfo.setImportPackages(config.getSuperEntityClass());
            } else {
                tableInfo.setImportPackages(Serializable.class.getCanonicalName());
            }
            // Boolean类型is前缀处理
            if (config.getStrategyConfig().isEntityBooleanColumnRemoveIsPrefix()) {
                tableInfo.getFields().stream().filter(field -> "boolean".equalsIgnoreCase(field.getPropertyType()))
                        .filter(field -> field.getPropertyName().startsWith("is"))
                        .forEach(field -> field.setPropertyName(config.getStrategyConfig(),
                                StringUtils.removePrefixAfterPrefixToLower(field.getPropertyName(), 2)));
            }
            tableInfo.setMapperName(tableInfo.getMapperName().replace("Mapper","Repository"));
            for (TableField field : tableInfo.getFields()) {
                String propertyName = field.getName();
                String comment = field.getComment();
                if (field.isKeyFlag()){

                    tableInfo.setImportPackages(Id.class.getCanonicalName());

                }
                Optional.ofNullable(fieldTypeHandler)
                        .map(handler->handler.apply(tableInfo.getName(),propertyName))
                        .ifPresent(columnType->{
                            tableInfo.setImportPackages(columnType.getPkg());
                            field.setColumnType(columnType);
                            if (columnType instanceof IEnumColumnType){
                                Map<String, Object> map = Optional
                                        .ofNullable(field.getCustomMap())
                                        .orElseGet(()->{
                                            field.setCustomMap(new HashMap<>());
                                            return field.getCustomMap();
                                        });

                                if (((IEnumColumnType) columnType).isEnum()){
                                    if (field.getPropertyType().contains("String")){

                                        map.put("enumType","String");
                                    }else{
                                        map.put("enumType","ORDINAL");
                                    }

                                }

                            }
                        });

                Optional.ofNullable(fieldValueGenerationHandler)
                        .map(handler->handler.apply(tableInfo.getName(),propertyName))
                        .ifPresent(iColumnType -> {
                            tableInfo.setImportPackages(iColumnType.getPkg());

                            field.setFill(iColumnType.getType());
                        });

            }


        }
        return config.setTableInfoList(tableList);

    }


    public void setFieldTypeHandler(BiFunction<String, String, IColumnType> fieldTypeHandler) {
        this.fieldTypeHandler = fieldTypeHandler;
    }

    public void setFieldValueGenerationHandler(BiFunction<String, String, IColumnType> fieldValueGenerationHandler) {
        this.fieldValueGenerationHandler = fieldValueGenerationHandler;
    }
}
